#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：arXiv-analyzer
@File    ：data_processor.py
@IDE     ：PyCharm
@Author  ：iyoahs
@Date    ：2025/6/20 21:45
@Describe：
'''
from pymongo import MongoClient
from typing import List, Dict
from src.utils.logger import setup_logger
import datetime

logger = setup_logger()


class DataProcessor:
    def __init__(self, mongo_config: Dict):
        self.mongo_config = mongo_config
        self.uri = mongo_config["uri"]
        self.db_name = mongo_config["db_name"]
        self.collection_name = mongo_config["collection_name"]
        self.processed_collection_name = mongo_config["processed_collection_name"]

    def clean_text(self, text: str) -> str:
        """简单的文本清洗"""
        return text.replace('\n', ' ').strip()

    def process_papers(self) -> List[Dict]:
        """处理论文数据"""
        try:
            client = MongoClient(self.uri)
            db = client[self.db_name]
            input_col = db[self.collection_name]
            output_col = db[self.processed_collection_name]

            papers = list(input_col.find())
            processed_papers = []

            for paper in papers:
                processed = {
                    "id": paper["id"],
                    "title": self.clean_text(paper["title"]),
                    "abstract": self.clean_text(paper["abstract"]),
                    "authors": paper["authors"],
                    "published": paper["published"],
                    "url": paper["url"],
                    "processed_at": datetime.datetime.utcnow().isoformat()
                }
                processed_papers.append(processed)

            # 批量写入处理后的数据
            if processed_papers:
                output_col.insert_many(processed_papers)
                logger.info(f"Processed {len(processed_papers)} papers")

            return processed_papers
        except Exception as e:
            logger.error(f"Processing failed: {str(e)}")
            return []
        finally:
            client.close()